import h5py
import random
import numpy as np
#from mpi4py import MPI

class HDF5():
    def __init__(self,filePath,mode='r'):
        #self.f = h5py.File(filePath, mode,driver='mpio', comm=MPI.COMM_WORLD)
        self.f = h5py.File(filePath, mode)
        self.grps=list(self.f.keys())

    def add_group(self,grpName):    #向h5文件第一层中添加一个group
        return self.f.create_group(grpName)

    def add_dataset(self,grpName,dsetName,data):    #在group下添加一个dataset
         #self.f[grpName].create_dataset(dsetName, data=data,compression="lzf")
         self.f[grpName].create_dataset(dsetName, data=data)    #暂时不用压缩,压缩大概可以减少10%的存储使用

    def set_grp_attrs(self,grpName,attrs,data): #给group添加属性
        #设置group属性，l_data表示group下有多少dataset
        self.f[grpName].attrs[attrs]=data

    def get_grp_attrs(self,grpName,attrs):  #获取group属性
        return self.f[grpName].attrs[attrs]

    def get_grp_size(self): #获取有多少个group，在我们的数据架构中，这个值表示有多少局对局
        return len(self.f.keys())

    def set_dset_attrs(self,grpName,dsetName,attrs,data):   #给dataset设置属性
        self.f[grpName][dsetName].attrs[attrs]=data

    def get_dset_attrs(self,grpName,dsetName,attrs):    #设置dataset的属性
        return self.f[grpName][dsetName].attrs[attrs]

    def get_rand_grpName(self): #随机读取一个group
        return random.choice(self.grps)

    def get_board_trans(self,npdata,way):   #给numpy数字变形，旋转或者镜像
        if way =='clockwise_0':
            return npdata
        elif way =='clockwise_1':
            return np.rot90(npdata,k=1)
        elif way =='clockwise_2':
            return np.rot90(npdata,k=2)
        elif way =='clockwise_3':
            return np.rot90(npdata,k=3)
        elif way =='anti_clockwise_0':
            return np.flip(npdata,axis=0)
        elif way =='anti_clockwise_1':
            return np.rot90(np.flip(npdata,axis=0),k=1)
        elif way =='anti_clockwise_2':
            return np.rot90(np.flip(npdata,axis=0),k=2)
        elif way =='anti_clockwise_3':
            return np.rot90(np.flip(npdata,axis=0),k=3)
        else:
            return npdata

    def get_alpha_board_trans(self,npdata,way):   #给numpy数字变形，旋转或者镜像
        if way =='clockwise_0':
            return npdata
        elif way =='clockwise_1':
            return np.rot90(npdata,k=1,axes=(1,2))
        elif way =='clockwise_2':
            return np.rot90(npdata,k=2,axes=(1,2))
        elif way =='clockwise_3':
            return np.rot90(npdata,k=3,axes=(1,2))
        elif way =='anti_clockwise_0':
            return np.flip(npdata)
        elif way =='anti_clockwise_1':
            return np.rot90(np.flip(npdata),k=1,axes=(1,2))
        elif way =='anti_clockwise_2':
            return np.rot90(np.flip(npdata),k=2,axes=(1,2))
        elif way =='anti_clockwise_3':
            return np.rot90(np.flip(npdata),k=3,axes=(1,2))
        else:
            return npdata

    def get_rand_dset_from_grp(self,grpName,isTrans=True,way='rand'):   #从group里随机挑选一个dataset，可以选择dataset是否要变形
        index=random.choice(list(self.f[grpName].keys()))
        data=np.array(self.f[grpName][index])
        if isTrans:
            x=data[0,:]
            y=data[-1,:]
            x=x.reshape(int(np.sqrt(x.size)),int(np.sqrt(x.size)))
            y=y.reshape(int(np.sqrt(y.size)),int(np.sqrt(y.size)))
            if way=='rand':
                choices=['clockwise_0','clockwise_1','clockwise_2','clockwise_3','anti_clockwise_0','anti_clockwise_1','anti_clockwise_2','anti_clockwise_3']
                choice=random.choice(choices)
                x=self.get_board_trans(x,choice).flatten()
                y=self.get_board_trans(y,choice).flatten()
            else:
                x=self.get_board_trans(x,way).flatten()
                y=self.get_board_trans(y,way).flatten()
            data=np.vstack((x,y))
        return data

    def get_rand_dset(self):    #从h5文件里随机挑选一个dataset出来
        grpName=self.get_rand_grpName()
        return self.get_rand_dset_from_grp(grpName)

    def get_rand_alphago_dset_from_grp(self,grpName,isTrans=True,p_or_v='value'):
        index=random.choice(list(self.f[grpName].keys()))
        data=np.array(self.f[grpName][index])
        boardsize=self.get_grp_attrs(grpName,'size')
        winner=self.get_grp_attrs(grpName,'winner')
        player=self.get_dset_attrs(grpName,index,'player')
        if p_or_v=='policy':
            move=self.get_dset_attrs(grpName,index,'move')
            npdata=np.zeros((boardsize,boardsize))
            if winner==player:
                npdata[tuple(move)]=1
            else:
                npdata[tuple(move)]=-1
        else:
            move=self.get_dset_attrs(grpName,index,'move')
            npmove=np.zeros((boardsize,boardsize))
            npmove[tuple(move)]=player
            if winner==player:
                npdata=np.array([1])
            else:
                npdata=np.array([-1])
        if isTrans:
            choices=['clockwise_0','clockwise_1','clockwise_2','clockwise_3','anti_clockwise_0','anti_clockwise_1','anti_clockwise_2','anti_clockwise_3']
            choice=random.choice(choices)
            data=self.get_alpha_board_trans(data,choice)
            if p_or_v=='policy':
                npdata=self.get_board_trans(npdata,choice)
            else:
                npmove=self.get_board_trans(npmove,choice)
        if p_or_v=='policy':
            return data,npdata
        else:
            return data,npmove,npdata

    def get_rand_alphogo_dset(self,p_or_v):
        grpName=self.get_rand_grpName()
        return self.get_rand_alphago_dset_from_grp(grpName,p_or_v=p_or_v)

    def yeilds_alphago_data(self,dataSize,boardSize,p_or_v):
        i = 0
        p_or_v=str(p_or_v,'utf-8')
        while i<dataSize:  #我们随机取出数据，不限制一下数据集大小，就会无限制的取值
        #while True:
            if p_or_v=='policy':
                x,y=self.get_rand_alphogo_dset(p_or_v)
                x=x.reshape(x.shape[1],x.shape[2],-1)   #channel_first转channel_last,Tensorflow做后端只支持这个格式
                yield x,y.flatten()
            else:
                x1,x2,y=self.get_rand_alphogo_dset(p_or_v)
                x1=x1.reshape(x1.shape[1],x1.shape[2],-1)
                yield {"input1":x1,"input2":x2.flatten()},y.flatten()
            i+=1
    def get_ac_dset(self):
        x_set=None
        y_set=None
        for i in self.grps:
            size=int(self.get_grp_attrs(i,'size'))
            winner=int(self.get_grp_attrs(i,'winner'))
            for k in self.f[i]:
                npboard=np.array(self.f[i][k])
                player=np.array(int(self.get_dset_attrs(i,k,'player')))
                target=tuple(self.get_dset_attrs(i,k,'target'))
                value=np.array(self.get_dset_attrs(i,k,'value'))
                nptarget=np.zeros((size,size))
                if player==winner:
                    nptarget[target]=1-value
                else:
                    nptarget[target]=-1-value
                x0=npboard
                y0=nptarget
                x1=np.rot90(x0, 1)
                y1=np.rot90(y0, 1)
                x2=np.rot90(x0, 2)
                y2=np.rot90(y0, 2)
                x3=np.rot90(x0, 3)
                y3=np.rot90(y0, 3)
                x4=np.flip(x0)
                y4=np.flip(y0)
                x5=np.rot90(x4, 1)
                y5=np.rot90(y4, 1)
                x6=np.rot90(x4, 2)
                y6=np.rot90(y4, 2)
                x7=np.rot90(x4, 3)
                y7=np.rot90(y4, 3)
                x0=x0.reshape(1,-1)
                x1=x1.reshape(1,-1)
                x2=x2.reshape(1,-1)
                x3=x3.reshape(1,-1)
                x4=x4.reshape(1,-1)
                x5=x5.reshape(1,-1)
                x6=x6.reshape(1,-1)
                x7=x7.reshape(1,-1)
                y0=y0.reshape(1,-1)
                y1=y1.reshape(1,-1)
                y2=y2.reshape(1,-1)
                y3=y3.reshape(1,-1)
                y4=y4.reshape(1,-1)
                y5=y5.reshape(1,-1)
                y6=y6.reshape(1,-1)
                y7=y7.reshape(1,-1)
                x0=np.hstack((x0,np.array(player).reshape(1,1)))
                x1=np.hstack((x1,np.array(player).reshape(1,1)))
                x2=np.hstack((x2,np.array(player).reshape(1,1)))
                x3=np.hstack((x3,np.array(player).reshape(1,1)))
                x4=np.hstack((x4,np.array(player).reshape(1,1)))
                x5=np.hstack((x5,np.array(player).reshape(1,1)))
                x6=np.hstack((x6,np.array(player).reshape(1,1)))
                x7=np.hstack((x7,np.array(player).reshape(1,1)))
                y0=np.hstack((y0,np.array(1 if player==winner else -1).reshape(1,1)))
                y1=np.hstack((y1,np.array(1 if player==winner else -1).reshape(1,1)))
                y2=np.hstack((y2,np.array(1 if player==winner else -1).reshape(1,1)))
                y3=np.hstack((y3,np.array(1 if player==winner else -1).reshape(1,1)))
                y4=np.hstack((y4,np.array(1 if player==winner else -1).reshape(1,1)))
                y5=np.hstack((y5,np.array(1 if player==winner else -1).reshape(1,1)))
                y6=np.hstack((y6,np.array(1 if player==winner else -1).reshape(1,1)))
                y7=np.hstack((y7,np.array(1 if player==winner else -1).reshape(1,1)))
                x=np.vstack((x0,x1,x2,x3,x4,x5,x6,x7))
                y=np.vstack((y0,y1,y2,y3,y4,y5,y6,y7))
                if x_set is None or y_set is None:
                    x_set=x
                    y_set=y
                else:
                    x_set=np.vstack((x_set,x))
                    y_set=np.vstack((y_set,y))
        return x_set,y_set

    def get_dl_dset(self):  #一次性取出全部数据
        x_set=None
        y_set=None
        for i in self.grps: #遍历group
            size=int(self.get_grp_attrs(i,'size'))
            winner=np.array(int(self.get_grp_attrs(i,'winner'))).reshape(1,1)
            '''
            data_box=np.zeros((1,3,size,size))  #数据量，通道，棋盘size，棋盘size
            if winner==1:
                d=np.ones((size,size))
            else:
                d=-1*np.ones((size,size))
            data_box[:,0,::]=d  #第0层存放该局对弈谁胜利的信息
            '''
            for k in self.f[i]: #遍历dset
                player=np.array(int(self.get_dset_attrs(i,k,'player'))).reshape(1,1)
                data=np.array(self.f[i][k])
                x=data[0,:].reshape(1,size,size)
                y=data[-1,:].reshape(1,size,size)
                z=np.vstack((x,y))  #合并在一起旋转，省得写两次
                z0=z
                #'''
                z1=np.rot90(z, 1, (1,2))
                z2=np.rot90(z, 2, (1,2))
                z3=np.rot90(z, 3, (1,2))
                z4=np.flip(z,axis=1)
                z5=np.rot90(np.flip(z,axis=1), 1, (1,2))
                z6=np.rot90(np.flip(z,axis=1), 2, (1,2))
                z7=np.rot90(np.flip(z,axis=1), 3, (1,2))
                #'''
                x0=z0[0].reshape(1,-1)
                y0=z0[1].reshape(1,-1)
                #'''
                x1=z1[0].reshape(1,-1)
                y1=z1[1].reshape(1,-1)
                x2=z2[0].reshape(1,-1)
                y2=z2[1].reshape(1,-1)
                x3=z3[0].reshape(1,-1)
                y3=z3[1].reshape(1,-1)
                x4=z4[0].reshape(1,-1)
                y4=z4[1].reshape(1,-1)
                x5=z5[0].reshape(1,-1)
                y5=z5[1].reshape(1,-1)
                x6=z6[0].reshape(1,-1)
                y6=z6[1].reshape(1,-1)
                x7=z7[0].reshape(1,-1)
                y7=z7[1].reshape(1,-1)
                #'''
                x0=np.hstack((x0,player,winner))
                #'''
                x1=np.hstack((x1,player,winner))
                x2=np.hstack((x2,player,winner))
                x3=np.hstack((x3,player,winner))
                x4=np.hstack((x4,player,winner))
                x5=np.hstack((x5,player,winner))
                x6=np.hstack((x6,player,winner))
                x7=np.hstack((x7,player,winner))


                x=np.vstack((x0,x1,x2,x3,x4,x5,x6,x7))
                y=np.vstack((y0,y1,y2,y3,y4,y5,y6,y7))
                #'''
                #x=x0
                #y=y0
                if x_set is None or y_set is None:
                    x_set=x
                    y_set=y
                else:
                    x_set=np.vstack((x_set,x))
                    y_set=np.vstack((y_set,y))
        return x_set,y_set


    def yeilds_data(self,dataSize,boardSize):
        i = 0
        while i<dataSize:  #我们随机取出数据，不限制一下数据集大小，就会无限制的取值
        #while True:
            npdata=self.get_rand_dset()
            yield np.concatenate((npdata[0,:].reshape(boardSize**2,),npdata[1,:].sum().reshape(1,).astype(int))),npdata[1,:].reshape(boardSize**2,)
            i+=1

    def isGrpExist(self,grpName):   #判断group名字是否已经在h5文件的第一层存在了
        if grpName in self.f.keys():
            return True
        else:
            return False
    def closeH5(self):  #关闭h5的句柄
        self.f.close()
